#include <cstdlib>
#include <iostream>
#include <string>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <poll.h>

#include "tcetools/socket.h"

#include "constants.h"
#include "io_loops.h"
#include "halfproxy_auth.h"
#include "daemon.h"

using namespace std;
using namespace tcetools;

const char *target = "localhost";
const char *target_port = NULL;
const char *bouncer = NULL;
const char *bouncer_port = NULL;
int tcp_wndsize = 0;

string secret("halfproxy");

void two_connects() {
	ClientSocket *s1, *s2;
	char buf[2];

	s1 = s2 = NULL;

	try {
		cout << "Child activated.\n"
			 << "Connecting to " << bouncer << ":" << bouncer_port << "..." << endl;
		s1 = new ClientSocket(bouncer, bouncer_port);
		cout << "Connected, now authenticating... " << flush;
		check_auth_client(*s1, secret, "DATA");
		if (s1->sck_closed()) {
			cout << "Failed" << endl;
			delete s1;
			return;
		} else cout << "OK" << endl;

		cout << "Connecting to " << target << ":" << target_port << "..." << endl;
		s2 = new ClientSocket(target, target_port);
		cout << "Connected to " << s2->get_rem_ip() << endl;

		IOLoop4(s1, s2, bouncer);
	} catch(SocketError &se) {
		cerr << "Error: " << se.what() << endl;
	}

	if (s1) delete s1;
	if (s2) delete s2;
}

void launch_proxy() {
	// Let's get back them status info to clear them zombie
	// processes
	int status, cleared = 0;
	while (waitpid(-1, &status, WNOHANG) > 0) cleared++;
	if (cleared) cout << "Cleared " << cleared << " zombie(s)." << endl;

	// Fork...
	int f_ret = fork();

	if (f_ret < 0) {
		cerr << "Warning: Unable to fork!" << endl;
	} else if (f_ret == 0) { // child
		two_connects();

		// BAI
		_exit(EXIT_SUCCESS);
	}
}

int main(int argc, char *argv[]) {
	int be_daemon = 0;
	const char *rdr_file = NULL;
	// Parse command line arguements
	if (argc > 0) {
		for (int i = 1; i < argc; i += 2) {
			if (string(argv[i]).compare("-sa") == 0) {
				if (i + 1 >= argc) {
					// Quit and print usage
					bouncer = NULL;
					break;
				}
				target = argv[i + 1];
			} else if (string(argv[i]).compare("-sp") == 0) {
				if (i + 1 >= argc) {
					// Quit and print usage
					bouncer = NULL;
					break;
				}
				target_port = argv[i + 1];
			} else if (string(argv[i]).compare("-ba") == 0) {
				if (i + 1 >= argc) {
					// Quit and print usage
					bouncer = NULL;
					break;
				}
				bouncer = argv[i + 1];
			} else if (string(argv[i]).compare("-bp") == 0) {
				if (i + 1 >= argc) {
					// Quit and print usage
					bouncer = NULL;
					break;
				}
				bouncer_port = argv[i + 1];
			} else if (string(argv[i]).compare("--secret") == 0) {
				if (i + 1 >= argc) {
					// Quit and print usage
					bouncer = NULL;
					break;
				}
				secret = argv[i + 1];
			} else if (string(argv[i]).compare("-w") == 0) {
				if (i + 1 >= argc) {
					// Quit and print usage
					bouncer = NULL;
					break;
				}
				tcp_wndsize = atoi(argv[i + 1]);
			} else if (string(argv[i]).compare("-D") == 0) {
				be_daemon = 1;
				i--; // No parameter
			} else if (string(argv[i]).compare("-L") == 0) {
				if (i + 1 >= argc) {
					// Quit and print usage
					bouncer = NULL;
					break;
				}
				rdr_file = argv[i + 1];
			} else {
				// Quit and print usage
				bouncer = NULL;
				break;
			}
		}
	}

	if (bouncer == NULL || bouncer_port == NULL || target_port == NULL) {
		cout << "You must specify the bouncer's address and port and the server port!" << endl;
		cout << "Usage: halfproxy OPTIONS\nList of options:\n"
			 << "\t-sa       <Server address, default: \"localhost\">\n"
			 << "\t-sp       <Server port, required>\n"
			 << "\t--secret  <secret, default: \"halfproxy\">\n"
			 << "\t-ba       <Connection bouncer address, required>\n"
			 << "\t-bp       <Connection bouncer port, required>\n"
			 << "\t-w        <Set TCP window size>\n"
			 << "\t-D        Daemonize\n"
			 << "\t-L        <stderr and stdout redirect, used only with -D>\n"
			 << endl;
		return EXIT_FAILURE;
	}

	if (be_daemon) daemonize(rdr_file);

	cout << "Halfproxy (" << VERSION << ")\n"
		 << "----------------------------"
		 << "\nServer target is: " << target << ":" << target_port
		 << "\nConnection bouncer is: " << bouncer << ":" << bouncer_port
		 << "\nShared secret is: " << secret
		 << "\n" << endl;

	// Drop root
	if (geteuid() == 0) { // Hmm.... we are root...
		uid_t real_id = getuid();

		if (real_id == 0) {
			cout << "Current effective user is root, and the real ID is also root. "
					"It is more secure to chown the binary to root, and flip the setuid "
					"bit. Then I can change uid to the real (unprivilaged) "
					"user." << endl;
		} else {
			cout << "Current effective user is root. Changing effective uid to "
					<< real_id << "." << endl;
			seteuid(real_id);
		}
	}

	Socket::ignore_sigpipe();

	while(1) {
		try {
			cout << "Making control connection... " << flush;
			ClientSocket ctrl(bouncer, bouncer_port);
			cout << "OK" << endl;

			cout << "Starting authentication... " << flush;
			check_auth_client(ctrl, secret, "CONTROL");
			if (ctrl.sck_closed()) {
				cout << "Failed" << endl;
				cout << "Waiting 10 seconds... " << flush;
				sleep(10);
				cout << "done" << endl;
				continue;
			} else cout << "OK" << endl;
			
			struct pollfd pfd;
			pfd.fd = ctrl.get_fd();
			pfd.events = POLLIN;
			pfd.revents = 0;

			int pret;

			while ((pret = poll(&pfd, 1, 15000)) >= 0) {
				if ((pfd.revents & POLLERR) ||
					(pfd.revents & POLLHUP) ||
					(pfd.revents & POLLNVAL)) {
					cout << "Control connection terminated." << endl;
					break;
				}

				if (pfd.revents & POLLIN) {
					char buf[2];
					ssize_t readlen = ctrl.raw_read(buf, 2);

					if (readlen == 0) {
						cout << "Control connection terminated." << endl;
						break;
					}

					if (buf[0] == 'G' && buf [1] == 'O') {
						launch_proxy();
					} else if (buf[0] == 'H' && buf [1] == 'I') {
						// Keel alive. Do nothing.
					} else {
						cout << "Unknown command on control channel!" << endl;
						break;
					}
					
				} else if (pret == 0) { // Send Keepalive
					ctrl.raw_write("HI", 2);
					continue;
				} else {
					cout << "poll() woke up, but nothing happened!" << endl;
					break;
				}


				pfd.revents = 0;
			}

			cout << "Closing shop..." << endl;
			ctrl.close_sock();

		} catch (SocketError &se) {
			// Silently ignore...
			cout << "\nCaught a socket error: " << se.what() << endl;
		}

		cout << "Waiting 10 seconds... " << flush;
		sleep(10);
		cout << "done" << endl;
	}
}
